(sry, I am leaving for vacation and lazy to open PR).
Since you are refactoring MOE, here is another trick or bug we found. When we use DP2EP, the gradient reduction denominator (of EP) is always smaller than other modules(e.g. attention or embedding) so the actual gradient of EP is always higher. See logs.

we fixed this by adding a loss_average_denominator in parallel_dims.py and forcing EP's reduce denominator to

            transformer_block.feed_forward.experts.set_reduce_scatter_divide_factor(
                loss_average_denominator,
            )

in apply_fsdp, you can also check full code here

where loss_average_denominator = dp_replicate * dp_shard *cp see here

We also have a better version of the bias update that only needs one reduce, you can check the code here. A tricks here does not affect to Bias update, but we need to know, once we have activation checkpoint.

        if self.load_balance_coeff is not None:
           with torch.no_grad():
               self.tokens_per_expert.add_(num_tokens_per_expert)

It will be called more than once, so the actual stats value of num_tokens_per_expert will be 2x.

Nice @rakkit, we found the same issue with the ep grads being off by a factor. I was finding that set_reduce_scatter_divide_factor errored when using an mp policy, though.

Surprised you didn't hit that? Think I saw you're on torch==2.6 in another comment elsewhere

lol @garrett361 thanks for the info. I did not see the issue on both Torch 2.6 and 2.7.1.

To clarify I only test the default mp set(mixed_precision_param=bf16 and mixed_precision_reduce=fp32)